package com.pine.common.threadpool;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadPoolExecutor;

/**
 * 拒绝策略工厂
 *
 * @author pine manage
 * @since 2024-08-28
 */
public class RejectedPolicyFactory {

    private static final Map<String, RejectedExecutionHandler> REJECTED_POLICY_MAP = new HashMap<>(8);

    static {
        initRejectedPolicyMap();
    }

    /**
     * 注册拒绝策略
     *
     * @param rejectedPolicyCode       拒绝策略编码
     * @param rejectedExecutionHandler 拒绝策略实例
     */
    public static void registerRejectedPolicy(String rejectedPolicyCode, RejectedExecutionHandler rejectedExecutionHandler) {
        REJECTED_POLICY_MAP.put(rejectedPolicyCode, rejectedExecutionHandler);
    }

    /**
     * 获取拒绝策略，当传入的拒绝策略编码不存在时，默认使用CallerRunsPolicy
     *
     * @param rejectedPolicyCode 拒绝策略编码
     * @return 拒绝策略实例
     */
    public static RejectedExecutionHandler getRejectedPolicy(String rejectedPolicyCode) {
        if (rejectedPolicyCode == null) {
            return new ThreadPoolExecutor.CallerRunsPolicy();
        }

        return REJECTED_POLICY_MAP.getOrDefault(rejectedPolicyCode, new ThreadPoolExecutor.CallerRunsPolicy());
    }


    /**
     * 初始化拒绝策略
     */
    private static void initRejectedPolicyMap() {
        REJECTED_POLICY_MAP.put(RejectedPolicyCodeEnum.ABORT_POLICY.getCode(), new ThreadPoolExecutor.AbortPolicy());
        REJECTED_POLICY_MAP.put(RejectedPolicyCodeEnum.CALLER_RUNS_POLICY.getCode(), new ThreadPoolExecutor.CallerRunsPolicy());
        REJECTED_POLICY_MAP.put(RejectedPolicyCodeEnum.DISCARD_OLDEST_POLICY.getCode(), new ThreadPoolExecutor.DiscardOldestPolicy());
        REJECTED_POLICY_MAP.put(RejectedPolicyCodeEnum.DISCARD_POLICY.getCode(), new ThreadPoolExecutor.DiscardPolicy());
    }
}
